package main

import (
	"flag"
	"net/http"
	"net/url"
	"os"
	"strings"

	"github.com/sirupsen/logrus"
)

// isRegistryHost 判断字符串是否为域名或IP地址
func isRegistryHost(host string) bool {
	// 检查是否为域名（包含点号且符合域名格式）
	if strings.Contains(host, ".") || strings.Contains(host, ":") {
		return true
	}

	return false
}

func main() {
	debug := flag.Bool("d", false, "打印调试信息")
	printInfo := flag.Bool("pi", false, "只打印获取信息")
	arch := flag.String("a", "", "指定架构，不指定时可手动选择")
	proxy := flag.String("x", "", "http 代理地址")
	username := flag.String("u", "", "用户名")
	password := flag.String("p", "", "密码")
	tags := flag.Bool("t", false, "获取tag列表")
	syncCount := flag.Int("s", 3, "并发下载数量")
	var registry string
	flag.StringVar(&registry, "r", "registry-1.docker.io", "指定镜像仓库")

	flag.Parse()

	if *debug {
		SetLogLevel(logrus.InfoLevel)
	} else {
		SetLogLevel(logrus.ErrorLevel)
	}

	args := flag.Args()
	logrus.Debugln("输入参数为", args)

	if len(args) == 0 {
		logrus.Fatalln("请输入需要下载的包名")
	}

	var pkg = args[0]
	var tag string
	if len(args) > 1 {
		tag = args[1]
	} else {
		// 先检查是否为URL格式，避免URL中的冒号被误认为tag分隔符
		if strings.HasPrefix(pkg, "http://") || strings.HasPrefix(pkg, "https://") {
			// 处理URL格式
			parsedURL, err := url.Parse(pkg)
			if err == nil {
				// 保留完整的URL（包括协议）作为registry
				registry = parsedURL.Scheme + "://" + parsedURL.Host
				pkg = strings.TrimPrefix(parsedURL.Path, "/")
				if pkg == "" {
					logrus.Fatalln("URL中缺少镜像路径")
				}
			} else {
				logrus.Fatalln("无效的URL格式:", err)
			}
			// 检查pkg中是否包含tag
			var found bool
			pkg, tag, found = strings.Cut(pkg, ":")
			if !found {
				tag = "latest"
			}
		} else {
			// 原有的逻辑：处理普通镜像名
			// 先检查是否包含registry，再处理tag
			if registry == "registry-1.docker.io" {
				// 从镜像名中提取registry
				partsOfPkg := strings.Split(pkg, "/")
				if len(partsOfPkg) >= 2 && isRegistryHost(partsOfPkg[0]) {
					registry = partsOfPkg[0]
					pkg = strings.Join(partsOfPkg[1:], "/")
				}
			}
			// 然后处理tag
			var found bool
			pkg, tag, found = strings.Cut(pkg, ":")
			if !found {
				tag = "latest"
			}
		}
	}

	var client Client

	// 设置代理
	if *proxy != "" {
		proxyUrl, err := url.Parse(*proxy)
		if err != nil {
			logrus.Fatalln("代理地址"+*proxy+"错误", err)
		}
		logrus.Info("use http proxy ", proxyUrl.String())

		// 设置代理传输
		client.SetTransport(&http.Transport{
			Proxy: http.ProxyURL(proxyUrl),
		})
	}

	err := client.Pull(*syncCount, registry, pkg, tag, *arch, *printInfo, *tags, *username, *password)
	if err != nil {
		// 检查是否为认证错误
		if strings.Contains(err.Error(), "UNAUTHORIZED") || strings.Contains(err.Error(), "authentication required") {
			logrus.Fatalln("认证失败：该镜像需要登录认证，请使用 -u 和 -p 参数提供用户名和密码")
		}
		logrus.Fatalln("下载发生错误", err)
	}
	os.Exit(0)
}
